##########################################################################################

library(data.table)
library(optparse)
library(ggplot2)
library(Seurat)
library(monocle)
library(ArchR)
library(dplyr)

##########################################################################################
option_list <- list(
    make_option(c("--rna_data_monocle_file"), type = "character"),
    make_option(c("--comine_data_file"), type = "character"),
    make_option(c("--time_diff_gene_file"), type = "character"),
    make_option(c("--out_path"), type = "character") 
)

if(1!=1){
    
    ## 单细胞表达拟时序文件
    rna_data_monocle_file <- "~/20231121_singleMuti/results/monocole/testis.monocle.Rdata"

    ## 整合atac和rna的文件
    comine_data_file <- "~/20231121_singleMuti/results/atac_res/testis_combined_peak.combineRNA.Rdata"

    ## 差异表达基因文件
    time_diff_gene_file <- "~/20231121_singleMuti/results/monocole/pseudotime_ordergene.tsv"

    ## 输出
    out_path <- "~/20231121_singleMuti/results/monocole"

}

###########################################################################################
parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

rna_data_monocle_file <- opt$rna_data_monocle_file
comine_data_file <- opt$comine_data_file
time_diff_gene_file <- opt$time_diff_gene_file
out_path <- opt$out_path

dir.create(out_path , recursive = T)

###########################################################################################

a <- load(comine_data_file)
## testis_combined_peak_combineRNA

c <- load(rna_data_monocle_file)
## ALL_cds_all

## 逆时序差异基因列表
Time_diff <- data.frame(fread(time_diff_gene_file))

###########################################################################################
## RNA的时间
cell_time <- data.frame(cell_id = ALL_cds_all$cell , Pseudotime = ALL_cds_all$Pseudotime )
cell_time$cell_id <- gsub( '_' , '#' , cell_time$cell_id)
## 时间放到100之内,atac需要
cell_time$Pseudotime <- cell_time$Pseudotime/(max(cell_time$Pseudotime)) * 100

###########################################################################################
## RNA的拟时序结果
type <- "testis"

cell_level <- c("SSC","Differenting&Differented SPG",
    "Leptotene","Zygotene","Patchytene","Diplotene","Early stage of spermatids","Round&ElongateS.tids","Sperm"
    )
ALL_cds_all$cell_type <- factor(ALL_cds_all$cell_type , levels = cell_level)

## 黑色的点代码用于构造轨迹的差异基因；灰色是背景基因；
## 红色是根据计算的基因表达大小喝离散度分布的趋势
image_name <- paste0( out_path , "/plot_ordering_genes.",type,".pdf" )
plot2 <- plot_ordering_genes(ALL_cds_all)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)

## 拟时序的时间分布
image_name <- paste0( out_path , "/cell_type_trajectory",type,".pdf" )
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by="cell_type", cell_size=1)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)
image_name <- paste0( out_path , "/cell_type_trajectory",type,".divide.pdf" )
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by = "cell_type") + facet_wrap('~cell_type', nrow = 1)
ggsave(file = image_name , plot = plot2,width = 15,height = 5)

image_name <- paste0( out_path , "/clusters_trajectory",type,".pdf" )
ALL_cds_all$seurat_clusters_character <- as.character(ALL_cds_all$seurat_clusters)
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by="seurat_clusters_character", cell_size=1)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)
image_name <- paste0( out_path , "/clusters_trajectory",type,".divide.pdf" )
ALL_cds_all$seurat_clusters_character <- as.character(ALL_cds_all$seurat_clusters)
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by = "seurat_clusters_character") + facet_wrap('~seurat_clusters_character', nrow = 1)
ggsave(file = image_name , plot = plot2,width = 15,height = 5)

## 时序
image_name <- paste0( out_path , "/pseudotime_trajectory",type,".pdf" )
plot2 <- plot_cell_trajectory(ALL_cds_all, color_by="Pseudotime", cell_size=1)
ggsave(file = image_name , plot = plot2,width = 6.5,height = 6)


################################################################################
#### 与精子发生相关的重要TF
## Foxo1和Dmrt1是在未分化精原细胞（SSC）里面重要的

## E2f4是在分化的精原细胞里面活性较高的

## 早期还有Sohlh2，WT1

################################################################################
## 整合atac+rna的拟时序
projHeme5 <- testis_combined_peak_combineRNA
projHeme5@cellColData[cell_time$cell_id,"Pseudotime"] <- cell_time$Pseudotime

## atac基因活性的热图

trajGSM <- getTrajectory(ArchRProj = projHeme5, name = "Pseudotime", useMatrix = "GeneScoreMatrix", log2Norm = TRUE)
trajMM <- getTrajectory(ArchRProj = projHeme5, name = "Pseudotime", useMatrix = "GeneExpressionMatrix", log2Norm = FALSE)

## 去除基因名上的chr信息
rownames(trajGSM) <- sapply(strsplit( rownames(trajGSM) , ":" ) , "[" , 2)
rownames(trajMM) <- sapply(strsplit( rownames(trajMM) , ":" ) , "[" , 2)

## Integrative pseudo-time analyses
## 鉴定在时序过程中，RNA和ATAC改变趋势的基因
corGSM_MM <- correlateTrajectories(trajGSM, trajMM)
corGSM_MM[[1]]$name1 <- sapply(strsplit( corGSM_MM[[1]]$name1 , ":" ) , "[" , 2)
corGSM_MM[[1]]$name2 <- sapply(strsplit( corGSM_MM[[1]]$name2 , ":" ) , "[" , 2)

## 提取对应的基因
trajGSM2 <- trajGSM[corGSM_MM[[1]]$name1, ]
trajMM2 <- trajMM[corGSM_MM[[1]]$name2, ]

## 提取感兴趣的基因
#gene_use <- c("FOXO1" , "DMRT1", "E2F4" , "SOHLH2" , "WT1")
#trajGSM2 <- trajGSM[gene_use, ]
#trajMM2 <- trajMM[gene_use, ]

trajCombined <- trajGSM2
assay(trajCombined , withDimnames=FALSE) <- t(apply(assay(trajGSM2), 1, scale)) + t(apply(assay(trajMM2), 1, scale))
combinedMat <- plotTrajectoryHeatmap(trajCombined, returnMat = TRUE, varCutOff = 0)

rowOrder <- match(rownames(combinedMat), rownames(trajGSM2))
ht1 <- plotTrajectoryHeatmap(trajGSM2,  pal = paletteContinuous(set = "horizonExtra"),  varCutOff = 0, rowOrder = rowOrder)
ht2 <- plotTrajectoryHeatmap(trajMM2, pal = paletteContinuous(set = "solarExtra"), varCutOff = 0, rowOrder = rowOrder)

image_name <- paste0( out_path , "/plotTrajectoryHeatmap.atac.pdf")
pdf(image_name , width = 10, height = 8)
ht1 + ht2
dev.off()

## 该命令是画在archR的默认文件夹下面
#plotPDF(ht1 + ht2, name = "plotTrajectoryHeatmap.atac.pdf", ArchRProj = projHeme5, addDOC = FALSE, width = 10, height = 8)